# -*- coding: utf-8 -*-
import os
import sys
import warnings
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# 将自定义 pymagic 路径加入环境变量
pymagic_path = '/database/home/duansizhang/pywave'
if pymagic_path not in sys.path:
    sys.path.insert(0, pymagic_path)

from pymagic.tools.ecgdllinterface import DataSource, EcgAnalyzer
from hrvanalysis import (
    remove_outliers, interpolate_nan_values, remove_ectopic_beats,
    get_time_domain_features, get_frequency_domain_features,
    get_geometrical_features, get_sampen
)
import EntropyHub

warnings.filterwarnings("ignore")

# 1. 设置路径和参数
raw_file = '/database/private/mgcdb/raw/0a0b34b8220644848209d7a199afd6fc_1699056000000.raw'
out_dir  = '/database/home/duansizhang/hrv_predict/result'
os.makedirs(out_dir, exist_ok=True)

csv_global = os.path.join(out_dir, 'hrv_global_window.csv')
csv_rr     = os.path.join(out_dir, 'hrv_rr_features.csv')
img_ecg    = os.path.join(out_dir, 'ecg_segment.png')

fs      = 250     # 采样率
win_sec = 300     # 5分钟窗口

# 2. 读取 ECG 与 R 峰
hdata    = DataSource(raw_file, 3)
res      = EcgAnalyzer(hdata)
rpos_all = np.array(res['rpos'], dtype=int)
anntype  = np.array(res['anntype'], dtype=int)
filt_ecg = np.array(res['filterdata']).flatten()

# 3. 筛选 anntype == 1
rpos  = rpos_all[anntype == 1]
times = rpos / fs

# 4. 全局预处理一次
rr_ms     = np.diff(rpos) * 1000.0 / fs
rr_times  = times[1:]
rr_clean  = remove_outliers(rr_intervals=rr_ms, low_rri=300, high_rri=2000)
rr_interp = interpolate_nan_values(rr_intervals=rr_clean, interpolation_method='linear')
nn_beats  = remove_ectopic_beats(rr_intervals=rr_interp, method='malik')
nn_interp = interpolate_nan_values(rr_intervals=nn_beats)
mask_valid= ~np.isnan(nn_interp)
nn_array  = np.array(nn_interp)[mask_valid]
nn_times  = rr_times[mask_valid]

# 5. 特征提取函数（不含熵）
def extract_features(nn_seg):
    if len(nn_seg) < 3:
        return None
    feats = {}
    try:
        td = get_time_domain_features(nn_seg)
        feats.update({k: float(td[k]) if td[k] is not None else np.nan for k in td})
        fd = get_frequency_domain_features(nn_seg)
        feats.update({k: float(fd[k]) if fd[k] is not None else np.nan for k in fd})
        geo = get_geometrical_features(nn_seg)
        feats.update({k: float(geo[k]) if geo[k] is not None else np.nan for k in geo})
        feats['nbeats'] = len(nn_seg)
        feats['ApEn'] = np.nan
    except Exception:
        return None
    return feats

# 6. 构建 rows
rows = []

# 6.1 全局
g_feats = extract_features(nn_array)
if isinstance(g_feats, dict):
    rows.append({
        'label': 'global',
        'window_start': 0.0,
        'window_end': float(nn_times[-1]),
        'timestamp': 0.0,
        **g_feats
    })

# 6.2 滑窗（只取前 5 个窗口）
t_end = float(nn_times[-1])
window_starts = np.arange(0, t_end, win_sec)[:10]

for ws in window_starts:
    we = ws + win_sec
    idx = (nn_times >= ws) & (nn_times < we)
    seg = nn_array[idx]
    feats = extract_features(seg)
    if not isinstance(feats, dict):
        continue

    # —— 仅对滑窗段计算熵 —— 
    if len(seg) > 10:
        smen = get_sampen(seg)
        feats.update({k: float(smen[k]) if smen[k] is not None else np.nan for k in smen})
        try:
            apen_vals, _ = EntropyHub.ApEn(seg)
            feats['ApEn'] = float(apen_vals[2]) if len(apen_vals) > 2 else np.nan
        except AssertionError:
            feats['ApEn'] = np.nan

    rows.append({
        'label': 'window',
        'window_start': float(ws),
        'window_end': float(we),
        'timestamp': float(ws + win_sec/2),
        **feats
    })

# 7. 保存 CSV
df_all = pd.DataFrame(rows)
df_all.to_csv(csv_global, index=False, encoding='utf-8-sig')
print(f"已保存 全局+窗口 HRV -> {csv_global}")

# 8. 仅 RR 全局特征
if isinstance(g_feats, dict):
    pd.DataFrame([g_feats]).to_csv(csv_rr, index=False, encoding='utf-8-sig')
    print(f"已保存 仅 RR 特征 -> {csv_rr}")

# 9. ECG 可视化
seg_ecg = filt_ecg[:fs*10]
t_ecg   = np.arange(len(seg_ecg)) / fs
plt.figure(figsize=(10,3))
plt.plot(t_ecg, seg_ecg)
plt.title("ECG first 10s")
plt.xlabel("time(s)")
plt.ylabel("extent")
plt.grid(True)
plt.tight_layout()
plt.savefig(img_ecg, dpi=300)
plt.close()
print(f"已保存 ECG 图像 -> {img_ecg}")
